package com.xiaozhou.controller;

import com.xiaozhou.function.FunctionManager;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import com.zhipu.oapi.service.v4.model.params.CodeGeexExtra;
import com.zhipu.oapi.service.v4.model.params.CodeGeexTarget;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
 * @author shhzhou
 * @description: 代码模型 例子
 * @date 2024/11/4 9:55
 */
@RestController
@RequestMapping("/code")
public class CodeController {

    @Autowired
    private ClientV4 clientV4;

    @Autowired
    private FunctionManager functionManager;

    /**
     * @return
     * @return com.zhipu.oapi.service.v4.model.ModelApiResponse
     * @Author xiaozhou
     * @Description 代码问答 同步调用
     * @Date 11:06 2024/11/4
     * @Param message
     */
    @GetMapping("/syncModel")
    public ModelApiResponse syncModel(String message) {

        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("system", "你是一个代码专家，你叫小林，请根据用户的要求生成正确的代码。"));
        messages.add(new ChatMessage("user", message));

        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(false);
        request.setStop(stop);
        request.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);

        return modelApiResponse;
    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.model.ModelData>
     * @Author xiaozhou
     * @Description 代码问答 sse流式调用
     * @Date 11:27 2024/11/4
     * @Param message
     */
    @GetMapping("/sseModel")
    public Flux<ModelData> sseModel(String message) {

        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("system", "你是一个代码专家，你叫小林，请根据用户的要求生成正确的代码。"));
        messages.add(new ChatMessage("user", message));

        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(true);
        request.setStop(stop);
        request.setInvokeMethod(Constants.invokeMethodSse);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return modelData;
            }).doOnNext(modelData -> {
                sink.next(modelData);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @return
     * @return com.zhipu.oapi.service.v4.model.ModelApiResponse
     * @Author xiaozhou
     * @Description 代码问答 异步调用      id = 621217294747395139164236694169633419
     * @Date 11:43 2024/11/4
     * @Param message
     */
    @GetMapping("/asyncModel")
    public ModelApiResponse asyncModel(String message) {

        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("system", "你是一个代码专家，你叫小林，请根据用户的要求生成正确的代码。"));
        messages.add(new ChatMessage("user", message));

        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(false);
        request.setStop(stop);
        request.setInvokeMethod(Constants.invokeMethodAsync);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);

        return modelApiResponse;
    }


    /**
     * @return
     * @return com.zhipu.oapi.service.v4.model.ModelApiResponse
     * @Author xiaozhou
     * @Description 代码回答  代码补全  同步调用 范例
     * @Date 14:10 2024/11/4
     * @Param message
     */
    @GetMapping("/syncCompleteModel")
    public ModelApiResponse syncCompleteModel(String message) {

        // 构造用户消息
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("user", message));

        // 设置大模型触发停止语句
        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        // 设置大模型代码补全相关参数
        CodeGeexExtra codeGeexExtra = new CodeGeexExtra();

        CodeGeexTarget target = new CodeGeexTarget();
        target.setPath("src/main/java/com/xiaozhou/controller/CodeController.java");
        target.setLanguage("Java");
        target.setCodePrefix("String apple = \"apple\";");
        target.setCodeSuffix("String result = apple + a; return result;");

        codeGeexExtra.setTarget(target);
        codeGeexExtra.setContexts(new ArrayList<>());

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(false);
        request.setStop(stop);
        request.setExtra(codeGeexExtra);
        request.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);
        return modelApiResponse;

    }


    /**
     * @Author xiaozhou
     * @Description 代码回答  代码补全  sse调用 范例
     * @Date 14:19 2024/11/4
     * @Param message
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.model.ModelData>
     */
    @GetMapping("/sseCompleteModel")
    public Flux<ModelData> sseCompleteModel(String message) {

        // 构造用户消息
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("user", message));

        // 设置大模型触发停止语句
        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        // 设置大模型代码补全相关参数
        CodeGeexExtra codeGeexExtra = new CodeGeexExtra();

        CodeGeexTarget target = new CodeGeexTarget();
        target.setPath("src/main/java/com/xiaozhou/controller/CodeController.java");
        target.setLanguage("Java");
        target.setCodePrefix("String apple = \"apple\";");
        target.setCodeSuffix("String result = apple + a; return result;");

        codeGeexExtra.setTarget(target);
        codeGeexExtra.setContexts(new ArrayList<>());

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(true);
        request.setStop(stop);
        request.setExtra(codeGeexExtra);
        request.setInvokeMethod(Constants.invokeMethodSse);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return modelData;
            }).doOnNext(modelData -> {
                sink.next(modelData);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @Author xiaozhou
     * @Description 代码回答  代码补全  异步调用 范例
     * @Date 14:27 2024/11/4
     * @Param message
     * @return
     * @return com.zhipu.oapi.service.v4.model.ModelApiResponse
     */
    @GetMapping("/asyncCompleteModel")
    public ModelApiResponse asyncCompleteModel(String message) {

        // 构造用户消息
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("user", message));

        // 设置大模型触发停止语句
        List<String> stop = new ArrayList<>();
        stop.add("<|endoftext|>");
        stop.add("<|user|>");
        stop.add("<|assistant|>");
        stop.add("<|observation|>");

        // 设置大模型代码补全相关参数
        CodeGeexExtra codeGeexExtra = new CodeGeexExtra();

        CodeGeexTarget target = new CodeGeexTarget();
        target.setPath("src/main/java/com/xiaozhou/controller/CodeController.java");
        target.setLanguage("Java");
        target.setCodePrefix("String apple = \"apple\";");
        target.setCodeSuffix("String result = apple + a; return result;");

        codeGeexExtra.setTarget(target);
        codeGeexExtra.setContexts(new ArrayList<>());

        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("codegeex-4");
        request.setMessages(messages);
        request.setStream(false);
        request.setStop(stop);
        request.setExtra(codeGeexExtra);
        request.setInvokeMethod(Constants.invokeMethodAsync);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);
        return modelApiResponse;

    }


    /**
     * @return
     * @return com.zhipu.oapi.service.v4.model.QueryModelResultResponse
     * @Author xiaozhou
     * @Description 代码回答/代码补全 异步调用结果查询
     * @Date 12:06 2024/11/4
     * @Param id
     */
    @GetMapping("/callAsyncResult")
    public QueryModelResultResponse callAsyncResult(String id) {

        QueryModelResultRequest request = new QueryModelResultRequest();
        request.setTaskId(id);

        QueryModelResultResponse response = clientV4.queryModelResult(request);

        return response;

    }


}
